import cv2
import os
import sys
import json
import numpy as np
from tqdm import tqdm
from skimage import graph, segmentation, color
from skimage.measure import regionprops
import shutil


class NumpySafeEncoder(json.JSONEncoder):
    """自定义JSON编码器处理NumPy类型"""

    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)


def rgba_to_hsl(r, g, b, a=255):
    """将RGB颜色转换为HSL格式（优化版）"""
    if not (a > 0 and a <= 255):
        return {"H": 0, "S": 0, "L": 0, "R": 0, "G": 0, "B": 0}

    r_n = r / 255.0
    g_n = g / 255.0
    b_n = b / 255.0
    max_c = max(r_n, g_n, b_n)
    min_c = min(r_n, g_n, b_n)
    delta = max_c - min_c

    # 计算亮度
    l_val = (max_c + min_c) / 2.0

    if delta == 0:
        h_val = s_val = 0
    else:
        # 计算饱和度
        s_val = delta / (1 - abs(2 * l_val - 1))
        # 计算色相
        if max_c == r_n:
            h_val = (g_n - b_n) / delta + (6 if g_n < b_n else 0)
        elif max_c == g_n:
            h_val = (b_n - r_n) / delta + 2
        else:
            h_val = (r_n - g_n) / delta + 4
        h_val /= 6

    # 返回HSL值并转换为Python原生类型
    return {
        "H": round(h_val * 360, 2),
        "S": round(s_val * 100, 2),
        "L": round(l_val * 100, 2)
    }


def color_segmentation(image_path, k=8, min_pixels=0, color_thresh=15):
    # 读取图像
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    h, w, _ = image.shape

    # 清理并重建输出目录
    if os.path.exists('./split'):
        shutil.rmtree('./split')
    os.makedirs('./split', exist_ok=True)

    # K-means聚类
    pixel_data = image.reshape((-1, 3)).astype(np.float32)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
    _, labels, centers = cv2.kmeans(
        pixel_data, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
    centers = np.uint8(centers)
    clustered = centers[labels.flatten()].reshape((h, w, 3))
    labels = labels.reshape((h, w)) + 1  # 标签从1开始

    # 构建区域邻接图并融合小色块
    if min_pixels > 0:
        rag = graph.rag_mean_color(image, labels)
        regions = regionprops(labels)
        small_regions = [r.label for r in regions if r.area < min_pixels]
        neighbor_map = {r.label: [] for r in regions}

        # 查找相邻色块
        for region in regions:
            if region.label in small_regions:
                y0, x0, y1, x1 = region.bbox
                boundary_labels = np.unique(labels[max(0, y0-1):min(h, y1+1),
                                                   max(0, x0-1):min(w, x1+1)])
                neighbor_map[region.label] = [
                    x for x in boundary_labels if x != region.label and x not in small_regions]

        # 融合到最相似的相邻色块
        for small_label in small_regions:
            if not neighbor_map[small_label]:
                continue
            small_color = rag.nodes[small_label]['mean color']
            min_diff = float('inf')
            best_neighbor = None
            for neighbor in neighbor_map[small_label]:
                neighbor_color = rag.nodes[neighbor]['mean color']
                diff = np.linalg.norm(small_color - neighbor_color)
                if diff < min_diff and diff < color_thresh:
                    min_diff = diff
                    best_neighbor = neighbor
            if best_neighbor:
                labels[labels == small_label] = best_neighbor

    output_segments = {}
    unique_labels = np.unique(labels)

    # 处理每个色块
    for label in tqdm(unique_labels, desc="Processing segments"):
        if label == 0:  # 跳过背景
            continue

        # 创建色块掩码
        mask = np.where(labels == label, 255, 0).astype(np.uint8)
        mean_color = cv2.mean(image, mask)[:3]

        # 创建带透明通道的色块图像
        segmented_img = np.zeros((h, w, 4), dtype=np.uint8)
        segmented_img[..., :3] = mean_color
        segmented_img[..., 3] = mask

        # 创建HSL标记图（使用向量化操作优化性能）
        hsl_img = np.zeros((h, w), dtype=np.float32)
        non_transparent = (segmented_img[..., 3] > 0)
        hsl_img[non_transparent] = 1  # 标记非透明区域

        # 获取第一个非透明像素的颜色
        if np.any(non_transparent):
            y, x = np.argwhere(non_transparent)[0]
            r, g, b, a = segmented_img[y, x]
            split_color = rgba_to_hsl(r, g, b, a)
        else:
            split_color = {"H": 0, "S": 0, "L": 0, "R": 0, "G": 0, "B": 0}

        # 保存结果
        json_filename = f'segment_{label:03d}'
        output_segments[json_filename] = {
            "offset": {"H": 0, "S": 0, "L": 0},
            "split_color": split_color,
            "data": hsl_img.tolist()  # 转换为Python列表
        }

    # 使用自定义编码器输出JSON
    print(json.dumps(output_segments, cls=NumpySafeEncoder))
    return output_segments


# 主入口
if __name__ == "__main__":
    conf = json.loads(sys.argv[1])
    color_segmentation(
        conf['path'],
        conf.get('k', 8),
        conf.get('min_px', 0),
        conf.get('color_thresh', 15)
    )
